#define _CRT_SECURE_NO_WARNINGS

#include "cmd.h"
#include "pe.h"
#include "int_util.h"
#include <wctype.h>
#include <stdlib.h>

/****************************************************************************\
 * Private Types
\****************************************************************************/

typedef struct tag_cmd_id_map
{
	str_wc       str;
	const cmd_id id;
} cmd_id_map;

typedef struct tag_tbl_name_id_map
{
	str_wc         name;
	const table_id id;
} tbl_name_id_map;

/****************************************************************************\
 * Private Variables
\****************************************************************************/

static str_wc s_cmd_help_str =
	L"Commands:\n"
	L"  help | h | ?           Show this help.\n"
	L"  exit | quit | q        Exit the environment.\n"
	L"  open <file_path>       Open PE file that pointed by <file_path>.\n"
	L"  refresh | reopen       Re-open the current PE file.\n"
	L"  file                   Show full path of the current opened PE file.\n"
	L"\n"
	L"  dd<range>              Show file data in range of <range>.\n"
	L"                         <range> indicates byte offset in file.\n"
	L"                         For <range> format, see 'Remarks' section below.\n"
	L"\n"
	L"  rva <address>          Show information of an RVA.\n"
	L"  offset <address>       Show information of an offset in file.\n"
	L"\n"
	L"  hd.dos                 Show dos header (IMAGE_DOS_HEADER).\n"
	L"  hd.sig                 Show PE header signature ('PE\\0\\0').\n"
	L"  hd.file                Show PE header (IMAGE_FILE_HEADER).\n"
	L"  hd.opt                 Show optional header (IMAGE_OPTIONAL_HEADER).\n"
	L"  hd.sects               Show all section headers.\n"
	L"  hd.nt                  Show NT headers (IMAGE_NT_HEADER, includes IMAGE_NT_SIGNATURE,\n"
	L"                         IMAGE_FILE_HEADER, and IMAGE_OPTIONAL_HEADER).\n"
	L"\n"
	L"  md.clih                Show CLI header (IMAGE_COR20_HEADER).\n"
	L"  md.sts                 Show metadata storage signature.\n"
	L"  md.sth                 Show metadata storage headers.\n"
	L"  md.shs                 Show metadata stream headers.\n"
	L"  md.hdrs                Show all metadata headers (includes metadata storage signature and headers,\n"
	L"                         as well as metadata stream headers).\n"
	L"\n"
	L"  md.sns                 Show strong name signature.\n"
	L"  md.res <range>         Show resources data in range of <range>.\n"
	L"                         <range> indicates records (rows) in the ManifestResource table.\n"
	L"  md.vtf <range>         Show VTableFixups. <range> indicates 0-based index of VTableFixup item.\n"
	L"\n"
	L"  mdt.hdr                Show metadata table header.\n"
	L"  mdt.l                  List metadata tables information.\n"
	L"  mdt.lc                 List metadata tables information with columns information.\n"
	L"  mdt.ct                 List coded tokens information.\n"
	L"  mdt.ctrt               List coded tokens information with referenced tables.\n"
	L"\n"
	L"  mdt.t <name> <range>   List data of table <name> in range of <range>.\n"
	L"  mdt.ti <tid> <range>   List data of table with index <tid> in range of <range>.\n"
	L"                         <range> indicates record (row) indeies.\n"
	L"                         For table name and index, make use of 'mdt.l' command.\n"
	L"                         For <range> format, see 'Remarks' section below.\n"
	L"\n"
	L"  md.s[<offset>]         Show #Strings[<offset>], <offset> is offset in #Strings stream.\n"
	L"  md.us[<offset>]        Show #US[<offset>], <offset> is offset in #US stream.\n"
	L"  md.g[<offset>]         Show #GUID[<offset>], <offset> is 1-based index in #GUID stream.\n"
	L"  md.b[<offset>]         Show #Blob[<offset>], <offset> is offset in #Blob stream.\n"
	L"\n"
	L"  il.h <rid>             Show IL method header and extra sections only.\n"
	L"  il.f <rid>             Show IL method header, extra sections and IL instructions.\n"
	L"                         <rid> is 1-based row index within the Method table.\n"
	L"Remarks:\n"
	L"  <range>                Set a range to displya listed data, in format of:\n"
	L"                           [<offset>]          - show data at <offset>\n"
	L"                           [<offset>:<length>] - show <length> data items start at <offset>\n"
	L"                           [<start>-<end>]     - show data items from offset <start> to offset <end>\n"
	L"                         For example,\n"
	L"                           dd[80h]             - show byte data at 80h of the file.\n"
	L"                           dd[80h:4]           - show 4 bytes data at 80h of the file.\n"
	L"                           dd[80h-83h]         - show 4 bytes data from 80h to 83h of the file.\n"
	L"                           mdt.t Method[1]     - show the 1st row of Method table.\n"
	L"                           mdt.t Method[1:3]   - show the 1st, 2nd and 3rd rows of Method table.\n"
	L"                           mdt.t Method[1-3]   - equivalence to the above.\n"
	L"\n"
	L"  Numbers                Numbers used as <offset>, <length>, <start> and <end>\n"
	L"                         could be in format of:\n"
	L"                               decimal: 1234\n"
	L"                              octonary: 01234 or o1234\n"
	L"                           hexadecimal: 0x1234ABCD or 1234ABCDh\n"
//	L"                         Within the number, ',' can be used for grouping, for example,\n"
//	L"                           1,234,567,890 (decimal) or\n"
//	L"                           0,123,346,789 (octonary because the initial 0) or\n"
//	L"                           0x1234,ABCD or ABCD,1234h (hexadecimal).\n"
	L"\n"
	;

static const cmd_id_map s_cmd_map[] =
{
	// Should be stored in order, binary search will be applied.
	{ L"?",        CMD_HELP           },
	{ L"dd",       CMD_DUMP_DATA      },
	{ L"exit",     CMD_EXIT           },
	{ L"file",     CMD_FILE           },
	{ L"h",        CMD_HELP           },
	{ L"hd.dos",   CMD_HD_DOS         },
	{ L"hd.file",  CMD_HD_FILE        },
	{ L"hd.nt",    CMD_HD_NT          },
	{ L"hd.opt",   CMD_HD_OPTIONAL    },
	{ L"hd.sects", CMD_HD_SECTIONS    },
	{ L"hd.sig",   CMD_HD_PE_SIG      },
	{ L"help",     CMD_HELP           },
	{ L"il.f",     CMD_IL_FULL        },
	{ L"il.h",     CMD_IL_HDR         },
	{ L"md.b",     CMD_MD_BLOB        },
	{ L"md.clih",  CMD_MD_CLIH        },
	{ L"md.g",     CMD_MD_GUID        },
	{ L"md.hdrs",  CMD_MD_HDRS        },
	{ L"md.res",   CMD_MD_RES         },
	{ L"md.s",     CMD_MD_STR         },
	{ L"md.shs",   CMD_MD_STREAM_HDRS },
	{ L"md.sns",   CMD_MD_SNSIG       },
	{ L"md.sth",   CMD_MD_ST_HDR      },
	{ L"md.sts",   CMD_MD_ST_SIG      },
	{ L"md.us",    CMD_MD_US          },
	{ L"md.vtf",   CMD_MD_VTFIXUP     },
	{ L"mdt.ct",   CMD_MDT_CT         },
	{ L"mdt.ctrt", CMD_MDT_CTRT       },
	{ L"mdt.hdr",  CMD_MDT_HDR        },
	{ L"mdt.l",    CMD_MDT_LIST       },
	{ L"mdt.lc",   CMD_MDT_LIST_WC    },
	{ L"mdt.t",    CMD_MDT_T          },
	{ L"mdt.ti",   CMD_MDT_TI         },
	{ L"offset",   CMD_OFFSET         },
	{ L"open",     CMD_OPEN           },
	{ L"q",        CMD_EXIT           },
	{ L"quit",     CMD_EXIT           },
	{ L"refresh",  CMD_REFRESH        },
	{ L"reopen",   CMD_REFRESH        },
	{ L"rva",      CMD_RVA            },
};

static const tbl_name_id_map s_tbl_map[NUMBER_OF_METADATA_TABLES] =
{
	// Should be stored by name in order, binary search will be applied.
	{ L"Assembly",               TBLID_Assembly               },
	{ L"AssemblyOS",             TBLID_AssemblyOS             },
	{ L"AssemblyProcessor",      TBLID_AssemblyProcessor      },
	{ L"AssemblyRef",            TBLID_AssemblyRef            },
	{ L"AssemblyRefOS",          TBLID_AssemblyRefOS          },
	{ L"AssemblyRefProcessor",   TBLID_AssemblyRefProcessor   },
	{ L"ClassLayout",            TBLID_ClassLayout            },
	{ L"Constant",               TBLID_Constant               },
	{ L"CustomAttribute",        TBLID_CustomAttribute        },
	{ L"DeclSecurity",           TBLID_DeclSecurity           },
	{ L"ENCLog",                 TBLID_ENCLog                 },
	{ L"ENCMap",                 TBLID_ENCMap                 },
	{ L"Event",                  TBLID_Event                  },
	{ L"EventMap",               TBLID_EventMap               },
	{ L"EventPtr",               TBLID_EventPtr               },
	{ L"ExportedType",           TBLID_ExportedType           },
	{ L"Field",                  TBLID_Field                  },
	{ L"FieldLayout",            TBLID_FieldLayout            },
	{ L"FieldMarshal",           TBLID_FieldMarshal           },
	{ L"FieldPtr",               TBLID_FieldPtr               },
	{ L"FieldRVA",               TBLID_FieldRVA               },
	{ L"File",                   TBLID_File                   },
	{ L"GenericParam",           TBLID_GenericParam           },
	{ L"GenericParamConstraint", TBLID_GenericParamConstraint },
	{ L"ImplMap",                TBLID_ImplMap                },
	{ L"InterfaceImpl",          TBLID_InterfaceImpl          },
	{ L"ManifestResource",       TBLID_ManifestResource       },
	{ L"MemberRef",              TBLID_MemberRef              },
	{ L"Method",                 TBLID_Method                 },
	{ L"MethodImpl",             TBLID_MethodImpl             },
	{ L"MethodPtr",              TBLID_MethodPtr              },
	{ L"MethodSemantics",        TBLID_MethodSemantics        },
	{ L"MethodSpec",             TBLID_MethodSpec             },
	{ L"Module",                 TBLID_Module                 },
	{ L"ModuleRef",              TBLID_ModuleRef              },
	{ L"NestedClass",            TBLID_NestedClass            },
	{ L"Param",                  TBLID_Param                  },
	{ L"ParamPtr",               TBLID_ParamPtr               },
	{ L"Property",               TBLID_Property               },
	{ L"PropertyMap",            TBLID_PropertyMap            },
	{ L"PropertyPtr",            TBLID_PropertyPtr            },
	{ L"StandAloneSig",          TBLID_StandAloneSig          },
	{ L"TypeDef",                TBLID_TypeDef                },
	{ L"TypeRef",                TBLID_TypeRef                },
	{ L"TypeSpec",               TBLID_TypeSpec               },
};

#define S_CMD_MAP_SZ (sizeof(s_cmd_map) / sizeof(s_cmd_map[0]))

static str_wc s_em_unknown_cmd = L"Unknown command.";

#define S_TBL_NAME_BUF_SZ 23  // 22 + 1
static wchar_t s_tbl_name_buf[S_TBL_NAME_BUF_SZ];

typedef int (* s_pf_cmp)(const void *, const void *);

/****************************************************************************\
 * Private Functions
\****************************************************************************/

/*
 * Check whether the given char is a valid cmd char.
 *
 * Parameters:
 *   c - The char to check.
 *
 * Returns: If c is valid char that can appear in a command, returns TRUE;
 *          otherwise FALSE.
 */
static BOOL s_is_cmd_char(wchar_t c)
{
	return (c == L'.') || (c == L'?')
		|| (c >= L'a' && c <= L'z')
		|| (c >= L'A' && c <= L'Z');
}

/*
 * Used to search cmd_id in s_cmd_map.
 */
static int s_cmp_cmd(str_wc cmd_str, cmd_id_map *p_map)
{
	return wcscmp(cmd_str, p_map->str);
}

/*
 * Used to search cmd_id in s_tbl_map.
 */
static int s_cmp_tbl(str_wc tbl_name, tbl_name_id_map *p_map)
{
	return wcscmp(tbl_name, p_map->name);
}

/*
 * Get the cmd_id part in a command.
 *
 * Parameters:
 *   cmd_str - Pointer to the command line.
 *   ret_cmd - The cmd object whose id will be set.
 *
 * Returns: Next part of the command line,
 *          i.e. pointer to the char right after the cmd_id part.
 */
static str_wc s_get_cmd_id(str_wc cmd_str, cmd *ret_cmd)
{
#define CMD_ID_BUF_LEN 20

	int len = 0;
	str_wc p_cmd = cmd_str;
	wchar_t cmd_id_buf[CMD_ID_BUF_LEN];
	cmd_id_map *p_cmd_map;

	// Skip whitespaces.
	while(*p_cmd != L'\0' && iswspace(*p_cmd))
		++p_cmd;

	// Get cmd id part from the cmd line.
	while(p_cmd[len] != L'\0' && s_is_cmd_char(p_cmd[len]) && len < CMD_ID_BUF_LEN - 1)
		++len;

	wcsncpy(cmd_id_buf, p_cmd, len);
	cmd_id_buf[len] = L'\0';

	if(len == 0)
	{
		ret_cmd->id = CMD_EMPTY;
		return p_cmd + len;
	}

	// Get the cmd id and fill in the cmd struct.
	p_cmd_map = (cmd_id_map *)bsearch(cmd_id_buf, s_cmd_map, S_CMD_MAP_SZ,
		sizeof(s_cmd_map[0]), (s_pf_cmp)s_cmp_cmd);
	if(p_cmd_map)
		ret_cmd->id = p_cmd_map->id;
	else
	{
		ret_cmd->id = CMD_ERROR;
		ret_cmd->err_msg = s_em_unknown_cmd;
	}

	return p_cmd + len;
}

/*
 * Get the str_arg part in a command.
 *
 * Parameters:
 *   cmd_str - Pointer to the args part within the command line.
 *   ret_cmd - The cmd object whose str_arg will be set.
 *
 * Returns: None.
 */
static BOOL s_get_cmd_str_arg(str_wc cmd_str, cmd *ret_cmd)
{
	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	ret_cmd->str_arg = cmd_str;

	return TRUE;
}

/*
 * Get the addr_arg part in a command.
 *
 * Parameters:
 *   cmd_str - Pointer to the args part within the command line.
 *   ret_cmd - The cmd object whose addr_arg will be set.
 *
 * Returns: None.
 */
static BOOL s_get_cmd_szt_arg(str_wc cmd_str, str_wc param_name, cmd *ret_cmd)
{
	size_t number_len, matched_number;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	number_len = 0;
	while(is_int_dig(cmd_str[number_len]))
		++number_len;

	if(number_len == 0)
	{
		wprintf(L"Missing %s.\n", param_name);
		return FALSE;
	}

	if(!parse_int(cmd_str, number_len, &matched_number))
	{
		wprintf(L"%s is invalid.\n", param_name);
		return FALSE;
	}

	ret_cmd->szt_arg = matched_number;
	return TRUE;
}

/*
 * Get the range parts in a command.
 *
 * Parameters:
 *     cmd_str - Pointer to the args part within the command line.
 *   match_len - Whether to match ':<len>]' part.
 *     ret_cmd - The cmd object whose range_start and range_len will be set.
 *
 * Returns: None.
 */
static BOOL s_get_cmd_range(str_wc cmd_str, BOOL match_len, cmd *ret_cmd)
{
	BOOL is_matching_len;
	size_t number_len, matched_number;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	// Match '['
	if(*cmd_str != L'[')
	{
		wprintf(L"Missing range start delimiter ('[').\n");
		return FALSE;
	}

	++cmd_str;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	// Match <offset>.
	number_len = 0;
	while(is_int_dig(cmd_str[number_len]))
		++number_len;

	if(number_len == 0)
	{
		wprintf(L"Missing <offset>.\n");
		return FALSE;
	}

	if(!parse_int(cmd_str, number_len, &matched_number))
	{
		wprintf(L"<offset> is invalid.\n");
		return FALSE;
	}

	ret_cmd->range_start = matched_number;

	cmd_str += number_len;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	// Match ']'  -- i.e. dd[<offset>]
	if(*cmd_str == L']')
	{
		ret_cmd->range_len = 1;
		return TRUE;
	}

	if(!match_len)
	{
		wprintf(L"Missing range end delimiter (']').");

		if(*cmd_str == L':')
			wprintf(L" The command doesn't support <length> part.");
		else if(*cmd_str == L'-')
			wprintf(L" The command doesn't support <end> part.");

		wprintf(L"\n");

		return FALSE;
	}

	// Match ':' or '-'.  -- i.e. dd[<offset>:<length>] or dd[<start>-<end>]
	if(*cmd_str == L':')
		is_matching_len = TRUE;
	else if(*cmd_str == L'-')
		is_matching_len = FALSE;
	else
	{
		wprintf(L"Missing range end delimiter (']').\n");
		return FALSE;
	}

	++cmd_str;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	// Matching <length> or <end>.
	number_len = 0;
	while(is_int_dig(cmd_str[number_len]))
		++number_len;

	if(number_len == 0)
	{
		wprintf(L"Missing <%s>.\n", is_matching_len ? L"length" : L"end");
		return FALSE;
	}

	if(!parse_int(cmd_str, number_len, &matched_number))
	{
		wprintf(L"<%s> is invalid.\n", is_matching_len ? L"length" : L"end");
		return FALSE;
	}

	if(is_matching_len)
		ret_cmd->range_len = matched_number;
	else
	{
		if(matched_number < ret_cmd->range_start)
		{
			wprintf(L"<end> is less than <start>.\n");
			return FALSE;
		}

		ret_cmd->range_len = matched_number - ret_cmd->range_start + 1;
	}

	cmd_str += number_len;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	// Match ']'
	if(*cmd_str != L']')
	{
		wprintf(L"Missing range end delimiter (']').\n");
		return FALSE;
	}

	return TRUE;
}

/*
 * For CMD_MDT_T, get table index, and save to tid_arg.
 */
static str_wc s_get_cmd_tbl_idx_byname(str_wc cmd_str, cmd *ret_cmd)
{
	size_t len;
	tbl_name_id_map *p_map;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	for(len = 0;
		*cmd_str >= L'a' && *cmd_str <= 'z' || *cmd_str >= 'A' && *cmd_str <= 'Z'
		&& len < S_TBL_NAME_BUF_SZ - 1;
		++len, ++cmd_str)
	{
		s_tbl_name_buf[len] = *cmd_str;
	}

	if(len == 0)
	{
		wprintf(L"Missing <name> part.\n");
		return NULL;
	}

	if(len >= S_TBL_NAME_BUF_SZ - 1)
	{
		wprintf(L"<name> part is too long.\n");
		return NULL;
	}

	s_tbl_name_buf[len] = L'\0';

	p_map = (tbl_name_id_map *)bsearch(s_tbl_name_buf, s_tbl_map, NUMBER_OF_METADATA_TABLES,
		sizeof(s_tbl_map[0]), (s_pf_cmp)s_cmp_tbl);

	if(!p_map)
	{
		wprintf(L"Unknown table name: %s\n", s_tbl_name_buf);
		return NULL;
	}

	ret_cmd->tid_arg = p_map->id;
	return cmd_str;
}

/*
 * For CMD_MDT_TI, get table index, and save to tid_arg.
 */
static str_wc s_get_cmd_tbl_idx_byid(str_wc cmd_str, cmd *ret_cmd)
{
	size_t number_len, matched_number;

	// Skip whitespaces.
	while(*cmd_str != L'\0' && iswspace(*cmd_str))
		++cmd_str;

	number_len = 0;
	while(is_int_dig(cmd_str[number_len]))
		++number_len;

	if(number_len == 0)
	{
		wprintf(L"Missing <tid> part.\n");
		return NULL;
	}

	if(!parse_int(cmd_str, number_len, &matched_number))
	{
		wprintf(L"<tid> is invalid.\n");
		return NULL;
	}

	if(matched_number >= NUMBER_OF_METADATA_TABLES)
	{
		wprintf(L"<tid> out of range.\n");
		return NULL;
	}

	ret_cmd->tid_arg = (table_id)matched_number;
	return cmd_str + number_len;
}

/****************************************************************************\
 * Public Functions
\****************************************************************************/

void show_cmd_help(void)
{
	wprintf(s_cmd_help_str);
}

void parse_cmd(str_wc cmd_line, cmd *ret_cmd)
{
	BOOL success = TRUE;
	str_wc next_part;

	next_part = s_get_cmd_id(cmd_line, ret_cmd);

	switch(ret_cmd->id)
	{
	case CMD_ERROR:
		return;
	case CMD_OPEN:
		success = s_get_cmd_str_arg(next_part, ret_cmd);
		if(!*ret_cmd->str_arg)
		{
			wprintf(L"Missing <file_path> part.\n");
			success = FALSE;
		}
		break;
	case CMD_DUMP_DATA:
	case CMD_MD_RES:
	case CMD_MD_VTFIXUP:
		success = s_get_cmd_range(next_part, TRUE, ret_cmd);
		break;
	case CMD_RVA:
	case CMD_OFFSET:
		success = s_get_cmd_szt_arg(next_part, L"<address>", ret_cmd);
		break;
	case CMD_MD_STR:
	case CMD_MD_US:
	case CMD_MD_GUID:
	case CMD_MD_BLOB:
		success = s_get_cmd_range(next_part, FALSE, ret_cmd);
		break;
	case CMD_MDT_T:
		if(!(next_part = s_get_cmd_tbl_idx_byname(next_part, ret_cmd)))
			success = FALSE;
		else
			success = s_get_cmd_range(next_part, TRUE, ret_cmd);
		break;
	case CMD_MDT_TI:
		if(!(next_part = s_get_cmd_tbl_idx_byid(next_part, ret_cmd)))
			success = FALSE;
		else
			success = s_get_cmd_range(next_part, TRUE, ret_cmd);
		break;
	case CMD_IL_FULL:
	case CMD_IL_HDR:
		success = s_get_cmd_szt_arg(next_part, L"<rid>", ret_cmd);
		break;
	default:
		break;
	}

	if(!success)
	{
		ret_cmd->id      = CMD_ERROR;
		ret_cmd->err_msg = NULL;  // Message has already been printed.
	}
}
